#include "TransformEvaluator.h"
#include "moremath.h"

// a zero matrix
const Mat4d mZeroMat = Mat4d(vl_0);

void TransformEvaluator::resetNumFrames(int _frames){
	for(int i=0; i<mTrans.size(); i++) {
		for(int j=0; j<mTrans[i]->size(); j++){
			if((*mTrans[i])[j]) delete (*mTrans[i])[j];
		}
		mTrans[i]->clear();
		delete mTrans[i];

	}
	mTrans.clear();

	for(int i=0; i<mFirstDeriv.size(); i++) {
		for(int j=0; j<mFirstDeriv[i]->size(); j++){
			for(int k=0; k<mFirstDeriv[i]->at(j).size(); k++){
				if((*mFirstDeriv[i])[j][k]) delete (*mFirstDeriv[i])[j][k];
			}
		}
		mFirstDeriv[i]->clear();
		delete mFirstDeriv[i];
	}
	for(int i=0; i<mSecondDeriv.size(); i++) {
		mSecondDeriv[i]->clear();
		//delete mSecondDeriv[i];
	}

	for(int i=0; i<mNumFrames; i++){
		for(int j=0; j<mModel->getNumBodyLinks(); j++){
			for(int k=0; k<mDofs.size(); k++)
				delete[] mSecondDerivMap[i][j][k];
			delete[] mFirstDerivMap[i][j];
			delete[] mSecondDerivMap[i][j];
		}
		delete[] mTransMap[i];
		delete[] mFirstDerivMap[i];
		delete[] mSecondDerivMap[i];
	}
	delete[] mTransMap;
	delete[] mFirstDerivMap;
	delete[] mSecondDerivMap;

	// -------------------------------------------------------------------
	int numDofs = mModel->getNumDofs();
	int numLinks = mModel->getNumBodyLinks();
	mNumFrames= _frames;
	// resize the matrices
	mTrans.resize(mNumFrames);
	mFirstDeriv.resize(mNumFrames);
	mSecondDeriv.resize(mNumFrames);
	mTransMap = new bool*[mNumFrames];
	mFirstDerivMap = new bool**[mNumFrames];
	mSecondDerivMap = new bool***[mNumFrames];

	for(int t=0; t<mNumFrames; t++){
		mTrans[t] = new vector<Mat4d *>;
		mTrans[t]->resize(numLinks);
		mTransMap[t] = new bool[numLinks];
		memset(mTransMap[t], false, numLinks*sizeof(bool));

		mFirstDeriv[t] = new vector<vector<Mat4d*> >;
		mFirstDeriv[t]->resize(numLinks);
		mFirstDerivMap[t] = new bool* [numLinks];

		mSecondDeriv[t] = new vector<vector<vector<Mat4d*> > >;
		mSecondDeriv[t]->resize(numLinks);
		mSecondDerivMap[t] = new bool** [numLinks];

		for(int i=0; i<numLinks; i++){
			//mTrans[t]->at(i)=vl_0;
			mTrans[t]->at(i)=NULL;
			mFirstDeriv[t]->at(i).resize(numDofs);
			mFirstDerivMap[t][i] = new bool[numDofs];
			memset(mFirstDerivMap[t][i],false, numDofs*sizeof(bool));
			mSecondDeriv[t]->at(i).resize(numDofs);
			mSecondDerivMap[t][i] = new bool*[numDofs];
			for(int j=0; j<numDofs; j++){
				mSecondDerivMap[t][i][j] = new bool[numDofs];
				memset(mSecondDerivMap[t][i][j], false, numDofs*sizeof(bool));
			}
			for(int j=0; j<numDofs; j++){
				(*mSecondDeriv[t])[i][j].resize(numDofs);
				for(int k=0; k<numDofs; k++){
					//(*mSecondDeriv[t])[i][j][k]=vl_0;
					(*mSecondDeriv[t])[i][j][k]=NULL;
				}
				//(*mFirstDeriv[t])[i][j]=vl_0;
				(*mFirstDeriv[t])[i][j]=NULL;
			}

		}
	}


}

TransformEvaluator::TransformEvaluator(Model *_model, int _numFrames, double *_timestep, vector<bool> *_activeDofMap, int _offset){
	mTimeStep = _timestep;
	// active map
	mActiveDofMap = _activeDofMap;
	// model
	mModel = _model;
	mCurrFrame = -1;
	mOffset = _offset;

	int numDofs = mModel->getNumDofs();
	int numLinks = mModel->getNumBodyLinks();
	mNumFrames= _numFrames;
	mDofs.clear();	// inited by ModelData/Solver

	// resize the matrices
	mTrans.resize(mNumFrames);
	mFirstDeriv.resize(mNumFrames);
	mSecondDeriv.resize(mNumFrames);
	mTransMap = new bool*[mNumFrames];
	mFirstDerivMap = new bool**[mNumFrames];
	mSecondDerivMap = new bool***[mNumFrames];

	for(int t=0; t<mNumFrames; t++){
		mTrans[t] = new vector<Mat4d *>;
		mTrans[t]->resize(numLinks);
		mTransMap[t] = new bool[numLinks];
		memset(mTransMap[t], false, numLinks*sizeof(bool));

		mFirstDeriv[t] = new vector<vector<Mat4d*> >;
		mFirstDeriv[t]->resize(numLinks);
		mFirstDerivMap[t] = new bool* [numLinks];

		mSecondDeriv[t] = new vector<vector<vector<Mat4d*> > >;
		mSecondDeriv[t]->resize(numLinks);
		mSecondDerivMap[t] = new bool** [numLinks];

		for(int i=0; i<numLinks; i++){
			//mTrans[t]->at(i)=vl_0;
			mTrans[t]->at(i)=NULL;
			mFirstDeriv[t]->at(i).resize(numDofs);
			mFirstDerivMap[t][i] = new bool[numDofs];
			memset(mFirstDerivMap[t][i],false, numDofs*sizeof(bool));
			mSecondDeriv[t]->at(i).resize(numDofs);
			mSecondDerivMap[t][i] = new bool*[numDofs];
			for(int j=0; j<numDofs; j++){
				mSecondDerivMap[t][i][j] = new bool[numDofs];
				memset(mSecondDerivMap[t][i][j], false, numDofs*sizeof(bool));
			}
			for(int j=0; j<numDofs; j++){
				(*mSecondDeriv[t])[i][j].resize(numDofs);
				for(int k=0; k<numDofs; k++){
					//(*mSecondDeriv[t])[i][j][k]=vl_0;
					(*mSecondDeriv[t])[i][j][k]=NULL;
				}
				//(*mFirstDeriv[t])[i][j]=vl_0;
				(*mFirstDeriv[t])[i][j]=NULL;
			}
		
		}
	}
}

Mat4d TransformEvaluator::getTrans(BodyLink *b, int frame){
	int bid = b->getModelID();

	// already computed
	if(mTransMap[frame-mOffset][bid]){
		return *((*mTrans[frame-mOffset])[bid]);
	}
	// need to compute
	else {
		loadDofs(frame);
		// get from parent and multiply local transform
		if(b->getParentJoint()->getBodyIn()!=NULL){
			Mat4d m = getTrans(b->getParentJoint()->getBodyIn(), frame);
			Mat4d l = b->getLocalTransform();

			//(*mTrans[frame-mOffset])[bid] = m*l;
			if((*mTrans[frame-mOffset])[bid]==NULL) (*mTrans[frame-mOffset])[bid] = new Mat4d;
			*((*mTrans[frame-mOffset])[bid]) = m*l;

			mTransMap[frame-mOffset][bid] = true;
			return *((*mTrans[frame-mOffset])[bid]);
		}
		// in root node - should have been computed by now
		else {
			mTransMap[frame-mOffset][bid] = true;

			//(*mTrans[frame-mOffset])[bid] = b->getLocalTransform();
			if((*mTrans[frame-mOffset])[bid]==NULL) (*mTrans[frame-mOffset])[bid] = new Mat4d;
			*((*mTrans[frame-mOffset])[bid]) = b->getLocalTransform();
			
			return *((*mTrans[frame-mOffset])[bid]);
		}
	}
}

Mat4d TransformEvaluator::getInvTrans(BodyLink *b, int frame){
	Mat4d transM = getTrans(b, frame);
	sub(transM, 3, 3) = trans(sub(transM, 3, 3));	// transpose the rotation component R-> R^T
	sub(transM, 0, 3, 3, 1) = -sub(transM, 3, 3)*sub(transM, 0, 3, 3, 1);	// change the translation component t -> -R^T*t
	return transM;
}

Mat4d TransformEvaluator::getFirstDeriv(BodyLink *b, Dof *q, int frame){
	int bid = b->getModelID();
	int qid = q->getModelID();

	// if b doesnt depend on q then return 0
	if(!b->dependsOnDof[q->getModelID()]){
		mFirstDerivMap[frame-mOffset][bid][qid] = true;

		//(*mFirstDeriv[frame-mOffset])[bid][qid] = Mat4d(vl_0);
		//return (*mFirstDeriv[frame-mOffset])[bid][qid];
		return mZeroMat;
	}

	// active map
	if(!(*mActiveDofMap)[qid]){
		mFirstDerivMap[frame-mOffset][bid][qid] = true;

		//(*mFirstDeriv[frame-mOffset])[bid][qid] = Mat4d(vl_0);
		//return (*mFirstDeriv[frame-mOffset])[bid][qid];
		return mZeroMat;
	}

	// already computed
	if(mFirstDerivMap[frame-mOffset][bid][qid]){
		return *((*mFirstDeriv[frame-mOffset])[bid][qid]);
	}
	// need to compute
	else {
		loadDofs(frame);
		Mat4d der = vl_1;
		// if q lies in the bodylink then have to compute here
		for(int i=b->getParentJoint()->getNumTransforms()-1; i>=0; i--){
			if(b->getParentJoint()->getTransform(i)->getModelID() == mModel->mTransformMap[q->getModelID()]){
				// compute now - q present in the bodylink dofs
				b->getParentJoint()->getTransform(i)->applyDerivativeMat(q, &der);
				q->setDerivFlag(true);
			}
			else b->getParentJoint()->getTransform(i)->applyTransformMat(&der);
		}
		if(q->getDerivFlag()){
			q->setDerivFlag(false);	// restore to false
			if(b->getParentJoint()->getBodyIn()!=NULL){
				der = getTrans(b->getParentJoint()->getBodyIn(), frame)*der;
			}
		}
		else {	// q not present in the bodylink dofs
			// get from parent and multiply local transform
			if(b->getParentJoint()->getBodyIn()!=NULL){
				der = getFirstDeriv(b->getParentJoint()->getBodyIn(), q,frame)*der;	// since der already has a local transform stored
			}
			// in root node - should have been computed by now
			else {
				printf("TransformEvaluator::getFirstDeriv - Should not have reached the root node without computing any deriv\n");
				der = mZeroMat;
			}
		}

		//(*mFirstDeriv[frame-mOffset])[bid][qid] =  der;
		if((*mFirstDeriv[frame-mOffset])[bid][qid]==NULL) (*mFirstDeriv[frame-mOffset])[bid][qid] = new Mat4d;
		*((*mFirstDeriv[frame-mOffset])[bid][qid]) =  der;

		mFirstDerivMap[frame-mOffset][bid][qid] = true;
		return *((*mFirstDeriv[frame-mOffset])[bid][qid]);
	}
}

// IMPORTANT - only m[b][q1][q2] element is stored where q1<q2. for q2<q1, return m[b][q2][q1]
Mat4d TransformEvaluator::getSecondDeriv(BodyLink *b, Dof *q1, Dof *q2, int frame){
	//	loadDofs(frame);
	int bid = b->getModelID();
	int q1id = q1->getModelID();
	int q2id = q2->getModelID();

	// if b doesnt depend on either q1 or q2 then return 0
	if(!b->dependsOnDof[q1id] || !b->dependsOnDof[q2id]){
		mSecondDerivMap[frame-mOffset][bid][q1id][q2id]=true;
		mSecondDerivMap[frame-mOffset][bid][q2id][q1id]=true;

		//(*mSecondDeriv[frame-mOffset])[bid][q1id][q2id] = Mat4d(vl_0);
		//(*mSecondDeriv[frame-mOffset])[bid][q2id][q1id] = (*mSecondDeriv[frame-mOffset])[bid][q1id][q2id];
		//return (*mSecondDeriv[frame-mOffset])[bid][q1id][q2id];
		return mZeroMat;
	}

	// active map
	if(!(*mActiveDofMap)[q1id] || !(*mActiveDofMap)[q2id]){
		mSecondDerivMap[frame-mOffset][bid][q1id][q2id]=true;
		mSecondDerivMap[frame-mOffset][bid][q2id][q1id]=true;

		//(*mSecondDeriv[frame-mOffset])[bid][q1id][q2id] = Mat4d(vl_0);
		//(*mSecondDeriv[frame-mOffset])[bid][q2id][q1id] = (*mSecondDeriv[frame-mOffset])[bid][q1id][q2id];
		//return (*mSecondDeriv[frame-mOffset])[bid][q1id][q2id];
		return mZeroMat;
	}

	// already computed
	if(mSecondDerivMap[frame-mOffset][bid][q1id][q2id]){
		//return (*mSecondDeriv[frame-mOffset])[bid][q1id][q2id];
		if(q1id < q2id) return *((*mSecondDeriv[frame-mOffset])[bid][q1id][q2id]);
		else return *((*mSecondDeriv[frame-mOffset])[bid][q2id][q1id]);
	}
	// need to compute
	else {
		loadDofs(frame);
		Mat4d der = vl_1;
		// if q1 or q2 lies in the bodylink then have to compute here
		for(int i=b->getParentJoint()->getNumTransforms()-1; i>=0; i--){
			if(b->getParentJoint()->getTransform(i)->getModelID() == mModel->mTransformMap[q1id]){
				// q1 present
				q1->setDerivFlag(true);
				if(b->getParentJoint()->getTransform(i)->getModelID() == mModel->mTransformMap[q2id]){
					// both q1 q2 present in transformation
					b->getParentJoint()->getTransform(i)->applyDerivativeMat(q1, q2, &der);
					q2->setDerivFlag(true);
				}
				else {	// only q1 present
					b->getParentJoint()->getTransform(i)->applyDerivativeMat(q1, &der);
				}
			}
			else if(b->getParentJoint()->getTransform(i)->getModelID() == mModel->mTransformMap[q2id]){
				b->getParentJoint()->getTransform(i)->applyDerivativeMat(q2, &der);
				q2->setDerivFlag(true);
			}
			else {
				b->getParentJoint()->getTransform(i)->applyTransformMat(&der);
			}
		}

		if(q1->getDerivFlag() && q2->getDerivFlag()){
			if(b->getParentJoint()->getBodyIn()!=NULL){
				der = getTrans(b->getParentJoint()->getBodyIn(), frame)*der;
			}
			q1->setDerivFlag(false);
			q2->setDerivFlag(false);
		}
		else if (q1->getDerivFlag()){
			if(b->getParentJoint()->getBodyIn()!=NULL){
				der = getFirstDeriv(b->getParentJoint()->getBodyIn(), q2, frame)*der;
			}
			q1->setDerivFlag(false);
		}
		else if (q2->getDerivFlag()){
			if(b->getParentJoint()->getBodyIn()!=NULL){
				der = getFirstDeriv(b->getParentJoint()->getBodyIn(), q1, frame)*der;
			}
			q2->setDerivFlag(false);
		}
		else {
			if(b->getParentJoint()->getBodyIn()!=NULL){
				der = getSecondDeriv(b->getParentJoint()->getBodyIn(), q1, q2, frame)*der;
			}
			else {
				printf("TransformEvaluator::getSecondDeriv - Should not have reached here as should have been computed wrt both dofs\n");
				der = Mat4d(vl_0);
			}
		}

		mSecondDerivMap[frame-mOffset][bid][q1id][q2id]=true;
		mSecondDerivMap[frame-mOffset][bid][q2id][q1id]=true;

		//(*mSecondDeriv[frame-mOffset])[bid][q1id][q2id] = der;
		//(*mSecondDeriv[frame-mOffset])[bid][q2id][q1id] = (*mSecondDeriv[frame-mOffset])[bid][q1id][q2id];
		//return (*mSecondDeriv[frame-mOffset])[bid][q1id][q2id];
		if(q1id<q2id){
			if((*mSecondDeriv[frame-mOffset])[bid][q1id][q2id]==NULL) (*mSecondDeriv[frame-mOffset])[bid][q1id][q2id] = new Mat4d;
			*((*mSecondDeriv[frame-mOffset])[bid][q1id][q2id]) = der;
			return *((*mSecondDeriv[frame-mOffset])[bid][q1id][q2id]); 
		}
		else {
			if((*mSecondDeriv[frame-mOffset])[bid][q2id][q1id]==NULL) (*mSecondDeriv[frame-mOffset])[bid][q2id][q1id] = new Mat4d;
			*((*mSecondDeriv[frame-mOffset])[bid][q2id][q1id]) = der;
			return *((*mSecondDeriv[frame-mOffset])[bid][q2id][q1id]); 
		}
	}
}

Mat4d TransformEvaluator::getTransDot(BodyLink *bi, int _frame){
	Mat4d ret(vl_0);
	for(int j=0; j<mModel->getNumDofs(); j++){
		Mat4d dWdq = getFirstDeriv(bi, mModel->getDof(j), _frame);
		double qDot = getVelocity(mModel->getDof(j)->getModelID(), _frame);
		ret+=dWdq*qDot;
	}
	return ret;
}

Mat4d TransformEvaluator::getTransDotPartialDof(BodyLink *bi, Dof *qk, int _frame){
	Mat4d ret(vl_0);
	for(int j=0; j<mModel->getNumDofs(); j++){
		Mat4d dWdqdq = getSecondDeriv(bi, mModel->getDof(j), qk, _frame);
		double qDot = getVelocity(mModel->getDof(j)->getModelID(), _frame);
		ret += dWdqdq*qDot;
	}
	return ret;
}

Mat4d TransformEvaluator::getTransDotPartialVel(BodyLink *bi, Dof *qk, int _frame){
	Mat4d dWdq = getFirstDeriv(bi, mModel->getDof(qk->getModelID()), _frame);
	return dWdq;
}

Mat4d TransformEvaluator::getTransDotDot(BodyLink *bi, int _frame){
	Mat4d ret(vl_0);
	for(int j=0; j<mModel->getNumDofs(); j++){
		if (!(*mActiveDofMap)[mModel->getDof(j)->getModelID()]) continue;
		Mat4d dWdq = getFirstDeriv(bi, mModel->getDof(j), _frame);
		double qDotDot = getAccel(mModel->getDof(j)->getModelID(), _frame);
		ret+=dWdq*qDotDot;
		double qjDot = getVelocity(mModel->getDof(j)->getModelID(), _frame);
		Mat4d ret2(vl_0);
		for(int k=0; k<mModel->getNumDofs(); k++){
			if (!(*mActiveDofMap)[mModel->getDof(k)->getModelID()]) continue;
			Mat4d dWdqdq = getSecondDeriv(bi, mModel->getDof(j), mModel->getDof(k), _frame);
			if(dWdqdq!=mZeroMat){
				double qkDot = getVelocity(mModel->getDof(k)->getModelID(), _frame);
				ret2+=dWdqdq*qkDot;
			}
		}
		ret+=ret2*qjDot;
	}
	return ret;
}

Mat4d TransformEvaluator::getTransDotDotApprox(BodyLink *bi, int _frame){
	return (getTransDot(bi, _frame) - getTransDot(bi, _frame-1))*getVelToSlopeDeriv();
}

void TransformEvaluator::computeTrans(BodyLink *b, int frame){
	getTrans(b, frame);

	// calculate recursively for the entire subtree
	for(int i=0; i<b->getNumChildJoints(); i++){
		computeTrans(b->getChildJoint(i)->getBodyOut(), frame);
	}
}

void TransformEvaluator::computeFirstDeriv(BodyLink *b, int frame){
	// iterate over dofs lying above
	BodyLink *temp=b;
	while(1){
		for(int i=0; i<temp->getParentJoint()->getNumDofs(); i++){
			Dof *qk=temp->getParentJoint()->getDof(i);
			if(qk->getModelID()!=-1){
				getFirstDeriv(b, qk, frame);
			}
		}
		if(temp->getParentJoint()->getBodyIn()!=NULL) temp=temp->getParentJoint()->getBodyIn();
		else break;
	}

	// calculate recursively for the entire subtree
	for(int i=0; i<b->getNumChildJoints(); i++){
		computeFirstDeriv(b->getChildJoint(i)->getBodyOut(), frame);
	}
}

void TransformEvaluator::computeSecondDeriv(BodyLink *b, int frame){
	// iterate over dofs lying above
	BodyLink *temp=b;
	while(1){
		for(int i=0; i<temp->getParentJoint()->getNumDofs(); i++){
			Dof *qj=temp->getParentJoint()->getDof(i);
			BodyLink *temp2=b;
			while(1){
				for(int j=0; j<temp2->getParentJoint()->getNumDofs(); j++){
					Dof *qk=temp2->getParentJoint()->getDof(j);
					if(qj->getModelID()!=-1 && qk->getModelID()!=-1){
						getSecondDeriv(b, qj, qk, frame);
					}
				}
				if(temp2->getParentJoint()->getBodyIn()!=NULL) temp2=temp2->getParentJoint()->getBodyIn();
				else break;
			}
		}
		if(temp->getParentJoint()->getBodyIn()!=NULL) temp=temp->getParentJoint()->getBodyIn();
		else break;
	}

	// calculate recursively for the entire subtree
	for(int i=0; i<b->getNumChildJoints(); i++){
		computeSecondDeriv(b->getChildJoint(i)->getBodyOut(), frame);
	}
}

void TransformEvaluator::cycleStructures(int _frameStart, int _frameEnd){
	// move the pointers by one frame
	vector<Mat4d*> * tempTrans = mTrans[_frameStart];
	vector<vector<Mat4d*> > * tempFirstDeriv = mFirstDeriv[_frameStart];
	vector<vector<vector<Mat4d*> > >* tempSecondDeriv = mSecondDeriv[_frameStart];
	bool* tempTransMap = mTransMap[_frameStart];
	bool** tempFirstDerivMap = mFirstDerivMap[_frameStart];
	bool*** tempSecondDerivMap = mSecondDerivMap[_frameStart];

	for(int i=_frameStart; i<_frameEnd; i++){
		mTrans[i] = mTrans[i+1];
		mFirstDeriv[i] = mFirstDeriv[i+1];
		mSecondDeriv[i] = mSecondDeriv[i+1];
		mTransMap[i] = mTransMap[i+1];
		mFirstDerivMap[i] = mFirstDerivMap[i+1];
		mSecondDerivMap[i] = mSecondDerivMap[i+1];
	}

	mTrans[_frameEnd] = tempTrans;
	mFirstDeriv[_frameEnd] = tempFirstDeriv;
	mSecondDeriv[_frameEnd] = tempSecondDeriv;
	mTransMap[_frameEnd] = tempTransMap;
	mFirstDerivMap[_frameEnd] = tempFirstDerivMap;
	mSecondDerivMap[_frameEnd] = tempSecondDerivMap;
}

TransformEvaluator::~TransformEvaluator(){

	for(int i=0; i<mTrans.size(); i++) {
		for(int j=0; j<mTrans[i]->size(); j++){
			if((*mTrans[i])[j]) delete (*mTrans[i])[j];
		}
		mTrans[i]->clear();
		delete mTrans[i];

	}
	mTrans.clear();

	for(int i=0; i<mFirstDeriv.size(); i++) {
		for(int j=0; j<mFirstDeriv[i]->size(); j++){
			for(int k=0; k<mFirstDeriv[i]->at(j).size(); k++){
				if((*mFirstDeriv[i])[j][k]) delete (*mFirstDeriv[i])[j][k];
			}
		}
		mFirstDeriv[i]->clear();
		delete mFirstDeriv[i];
	}
	for(int i=0; i<mSecondDeriv.size(); i++) {
		mSecondDeriv[i]->clear();
		//delete mSecondDeriv[i];
	}

	for(int i=0; i<mNumFrames; i++){
		for(int j=0; j<mModel->getNumBodyLinks(); j++){
			for(int k=0; k<mDofs.size(); k++)
				delete[] mSecondDerivMap[i][j][k];
			delete[] mFirstDerivMap[i][j];
			delete[] mSecondDerivMap[i][j];
		}
		delete[] mTransMap[i];
		delete[] mFirstDerivMap[i];
		delete[] mSecondDerivMap[i];
	}
	delete[] mTransMap;
	delete[] mFirstDerivMap;
	delete[] mSecondDerivMap;

	for(int i=0; i<mDofs.size(); i++){
		delete mDofs[i];
	}
	mDofs.clear();
}